Adding readManyByPartitionKey API#48801
Conversation
…to users/fabianm/readManyByPK
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
Adds a new readManyByPartitionKey API surface to the Java Cosmos SDK (sync + async) and wires it through the Spark connector to support PK-only reads (including partial HPK), with query-plan-based validation for custom queries.
Changes:
- Added public
readManyByPartitionKeyoverloads inCosmosAsyncContainer/CosmosContainerand an internalAsyncDocumentClient+RxDocumentClientImplimplementation that groups PKs by physical partition and issues per-range queries. - Introduced
ReadManyByPartitionKeyQueryHelperto compose PK filters into user-provided SQL and added a new config knob for per-partition batching. - Added Spark support (UDF + PK serialization/parsing helper + reader) and unit/integration tests for query composition and end-to-end behavior.
Reviewed changes
Copilot reviewed 18 out of 18 changed files in this pull request and generated 11 comments.
Show a summary per file
| File | Description |
|---|---|
| sdk/cosmos/docs/readManyByPartitionKey-design.md | Design doc describing the new API, query validation, and Spark integration approach. |
| sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java | Adds a helper method to fetch query plans through the gateway for validation. |
| sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java | Implements readManyByPartitionKey execution, validation, PK→range grouping, batching, and concurrency. |
| sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java | New helper to build SqlQuerySpec by appending PK filters and extracting table aliases. |
| sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java | Adds config/env accessors for max PKs per per-partition query batch. |
| sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java | Adds internal interface method for PK-only read-many. |
| sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java | Adds sync readManyByPartitionKey overloads. |
| sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java | Adds async readManyByPartitionKey overloads and wiring to internal client. |
| sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java | Unit tests for SQL generation, alias extraction, and WHERE detection. |
| sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java | Emulator integration tests for single PK + HPK, partial HPK, projections, and query validation. |
| sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKeyITest.scala | Spark integration test for reading by PKs and empty result behavior. |
| sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala | Unit tests for PK string serialization/parsing helpers. |
| sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala | Spark UDF to compute _partitionKeyIdentity values. |
| sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala | Spark partition reader that calls new SDK API and converts results to rows. |
| sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala | Spark reader that maps input rows to PKs and streams results via the partition reader. |
| sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala | Helper for PK serialization/parsing used by the UDF and data source. |
| sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala | Adds Spark entry point to read-many by partition key, including PK extraction logic. |
| sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConstants.scala | Adds _partitionKeyIdentity constant. |
Comments suppressed due to low confidence (1)
sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala:1
- The error message has mismatched parentheses/quoting (
classOf<SparkRowItem])) which makes it harder to read and search for. Suggest correcting it to a clean, unambiguous string (e.g.,classOf[SparkRowItem]) to improve diagnosability.
// Copyright (c) Microsoft Corporation. All rights reserved.
…s/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
…s/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
…ntation/RxDocumentClientImpl.java Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
…ntation/ReadManyByPartitionKeyQueryHelper.java Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
…to users/fabianm/readManyByPK
|
@sdkReviewAgent |
|
✅ Review complete (41:15) Posted 2 inline comment(s). Steps: ✓ context, correctness, cross-sdk, design, history, past-prs, synthesis, test-coverage |
|
/azp run java - cosmos - spark |
|
/azp run java - cosmos - tests |
|
Azure Pipelines successfully started running 1 pipeline(s). |
|
@sdkReviewAgent |
|
Azure Pipelines successfully started running 1 pipeline(s). |
|
✅ Review complete (33:55) Posted 2 inline comment(s). Steps: ✓ context, correctness, cross-sdk, design, history, past-prs, synthesis, test-coverage |
|
@sdkReviewAgent |
|
/azp run java - cosmos - tests |
|
/azp run java - cosmos - spark |
|
Azure Pipelines successfully started running 1 pipeline(s). |
1 similar comment
|
Azure Pipelines successfully started running 1 pipeline(s). |
|
✅ Review complete (32:34) Posted 2 inline comment(s). Steps: ✓ context, correctness, cross-sdk, design, history, past-prs, synthesis, test-coverage |
|
Live test failures independent of the changes in this PR |
…public The upstream ReadManyByPartitionKeyQueryPlanValidationTest (added in Azure#48801) constructs PartitionedQueryExecutionInfo directly from a test package. Upstream/main exposed this ctor as public; restore that visibility so the new test compiles after the merge. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…V2 endpoint (#47759) * Route Execute Stored Procedure requests to Thin Proxy endpoint. * Route Execute Stored Procedure and QueryPlan requests to Thin Proxy endpoint. * Ensure QueryPlan gets routed to Gateway Service Endpoint (in non-TC + GW Mode) cases. * Ensure QueryPlan gets routed to Gateway Service Endpoint (in non-TC + GW Mode) cases. * Ensure QueryPlan gets routed to Gateway Service Endpoint (in non-TC + GW Mode) cases. * Obtain List<Range<String>> from List<PartitionKeyInternal>. * Obtain List<Range<String>> from List<PartitionKeyInternal>. * Adding query + thin-client tests. * Fixing tests. * Fixing tests. * Fixing tests. * Addressing review comments. * Addressing review comments. * Refactor thin-client E2E tests based on operation type. * Refactor thin-client E2E tests based on operation type. * Refactor thin-client E2E tests based on operation type. * Refactor thin-client E2E tests based on operation type. * Refactor thin-client E2E tests based on operation type. * Refactor thin-client E2E tests based on operation type. * Refactor thin-client E2E tests based on operation type. * Refactor thin-client E2E tests based on operation type. * Refactor thin-client E2E tests based on operation type. * Add SupportedQueryFeatures and QueryVersion RNTBD request headers for QueryPlan proxy routing Add RNTBD token mappings for x-ms-cosmos-supported-query-features (0x002B) and x-ms-cosmos-query-version (0x002C) so the thin client proxy can read these values from the RNTBD body when processing QueryPlan requests. Previously these headers were only set as HTTP headers by QueryPlanRetriever and were lost when QueryPlan was routed through the proxy path, since ThinClientStoreModel serializes requests as RNTBD (not HTTP headers). IDs match server-side proxy definitions per ADO PR 1982503. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Add change feed tests for FeedRange.forFullRange and forLogicalPartition Add testThinClientChangeFeedFullRange covering FeedRange.forFullRange() across multiple partition keys, and testThinClientChangeFeedPartitionKey covering FeedRange.forLogicalPartition with exact doc count + PK validation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Add thin client E2E test matrix documentation for QueryPlan review Documents all 59 thin client E2E tests across query (50), point operations (3), change feed (3), and stored procedures (3) with SQL, query features covered, and known account-side blockers. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Add SupportedQueryFeatures and QueryVersion RNTBD request headers for QueryPlan proxy routing Add RNTBD token mappings for x-ms-cosmos-supported-query-features (0x00F0) and x-ms-cosmos-query-version (0x00F1) so the thin client proxy can read these values from the RNTBD body when processing QueryPlan requests. IDs are provisional (0x00F0, 0x00F1) — must be coordinated with server-side proxy team. See ADO PR 1982503 for the proxy-side design. Note: The design doc listed 0x002B/0x002C but those are already assigned to PartitionKey/PartitionKeyRangeId in the Java SDK. Using 0x00F0/0x00F1 to avoid ID collision until final server-side IDs are assigned. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix testGetCurrentDateTime flaky assertion, add AAD auth support, RNTBD instructions - Fix testGetCurrentDateTime: assert ISO 8601 format instead of exact match (gateway and proxy return slightly different timestamps) - Add DefaultAzureCredential support via COSMOS.USE_AAD_AUTH system property for accounts with disableLocalAuth=true - Add RNTBD class reference as .github/instructions/rntbd.instructions.md - Add pom.xml system properties for THINCLIENT_ENABLED, HTTP2_ENABLED, USE_AAD_AUTH - Add beforeSuiteReuse mode for degraded accounts Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Refactor thin client query tests for reliability - Switch baseline from Gateway V1 to Direct TCP to avoid JVM config interference (THINCLIENT_ENABLED/HTTP2_ENABLED affect Gateway V1) - Assert :10250 endpoint only on Gateway V2 results (not baseline) - Rename helpers: assertDirectAndThinClientMatch (was gateway) - Document seedTestData schema in Javadoc - Remove 'Expected to fail' comments (account has vector search enabled) - Clean up class/method Javadoc Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix container leaks and Direct TCP AAD auth in tests - Fix container leak: get container reference before createContainer() so finally block can always delete. Use safeDeleteContainer() helper. - Fix Direct TCP client: apply AAD credential via applyCredential() for accounts with disableLocalAuth=true. All 89 thin client tests pass (80 query + 3 change feed + 3 point ops + 3 sprocs). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Use bulkDelete in AfterClass for seeded docs cleanup Replace sequential deleteItem loop with executeBulkOperations for faster and more reliable cleanup of seeded test documents. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Address PR review feedback 1. Remove rntbd.instructions.md (not part of PR) 2. Move MAPPER variable to top of PartitionKeyInternalTest 3. Remove COSMOS.REUSE_DATABASE_ID (beforeSuiteReuse, afterSuite guard, pom.xml) 4. Keep System.setProperty in ThinClientQueryE2ETest (required — extends TestSuiteBase not ThinClientTestBase, property must be set before client build) 5. Fix RNTBD IDs to match server-side RntbdConstants.cs (ADO PR 1982503): SupportedQueryFeatures=0x00FF (String), QueryVersion=0x0100 (SmallString) The SmallString type was the root cause of 400 BadRequest — String uses 2-byte length prefix, SmallString uses 1-byte, causing token stream misparse. All 89 thin client tests pass. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Align vector/FTS/hybrid tests with existing patterns - Add euclidean VectorDistance variant to vector search test - Add document ID comparison to FTS and hybrid tests (was count-only) - Add testFullTextScoreRanking: ORDER BY RANK FullTextScore with exact order comparison between Direct and thin client - All assertions now validate both result size AND content parity 90/90 thin client tests pass. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Address review agent comments + fix broken link Code fixes: - PartitionKeyInternalHelper: materialize fieldNames() Iterator in error msg - QueryPlanRetriever: throw IllegalStateException when partitionKeyDef null in thin client mode (was silent fallthrough) - RxDocumentClientImpl: add parentheses around compound && in boolean exprs - ThinClientStoreModel: use != instead of !(==) - ThinClientTestBase: remove stale 'uncomment' comment - ThinClientQueryE2ETest: accept status 0 (transport rejection) for invalid query — proxy returns transport error, not HTTP 400 - THINCLIENT_TEST_MATRIX.md: remove broken link, update methodology to reflect Direct TCP baseline Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Address review agent comments (round 2) - RntbdOperationType: add QueryPlan case to fromId()/fromType() lookups - PartitionKeyInternalHelper: sanitize PII from error message (log node type instead of raw partition key values) - THINCLIENT_TEST_MATRIX.md: remove PR-specific references and known blockers - PartitionedQueryExecutionInfo: remove stale queryRanges from backing ObjectNode after EPK conversion to prevent data inconsistency - QueryPlanRetriever: add TODO for 3 missing query features vs .NET SDK (ListAndSetAggregate, CountIf, HybridSearchSkipOrderByRewrite) 90/90 thin client tests pass. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Refactor queryRanges deserialization — detect format from response, no ObjectNode mutation - PartitionedQueryExecutionInfo: getQueryRanges() now inspects the first range's 'min' field to detect format (string=EPK hex, array=PK-internal) and applies the correct deserialization path automatically. - Remove ObjectNode.remove('queryRanges') — no longer mutate the backing JSON. The format is detected lazily on first access. - 3-arg constructor now takes PartitionKeyDefinition (not pre-computed ranges) so conversion happens inside getQueryRanges(), keeping the DTO self-contained. - QueryPlanRetriever: pass partitionKeyDefinition to constructor instead of pre-converting ranges externally. 90/90 thin client tests pass. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Use agnostic QueryRangesFormat hint instead of PartitionKeyDefinition presence Introduce QueryRangesFormat enum (EPK_HEX_STRING, PARTITION_KEY_INTERNAL_ARRAY) as the deserialization hint. The hint guides which path to try first, but getQueryRanges() always validates by inspecting the actual min node type — if the hint doesn't match the wire format, it falls through to the other path. Naming is transport-agnostic (no Proxy/ThinClient/Gateway references). PartitionKeyDefinition is only stored when needed for EPK conversion, not used as a format signal. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix fetchQueryPlanForValidation caller after merge Pass null for the DocumentCollection parameter so the public fetchQueryPlanForValidation entry point compiles against the new fetchQueryPlan signature introduced during the merge with upstream/main. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Make PartitionedQueryExecutionInfo(ObjectNode, RequestTimeline) ctor public The upstream ReadManyByPartitionKeyQueryPlanValidationTest (added in #48801) constructs PartitionedQueryExecutionInfo directly from a test package. Upstream/main exposed this ctor as public; restore that visibility so the new test compiles after the merge. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * [Cosmos] Default thin-client to enabled and add HTTP/2 connectivity-probe gate Adds an EndpointOrchestrator that fans out POST /connectivity-probe to every thin-client regional endpoint after each topology refresh. SDK only routes data-plane traffic to thin-client (Gateway V2) when all regional probes succeed across N consecutive refresh cycles (configurable via COSMOS.THINCLIENT_PROBE_FAILURE_THRESHOLD, default 2); otherwise traffic falls back to Gateway V1 at the next refresh boundary. No mid-flight fallback. Caveats: - Probe wiring is skipped entirely for Direct mode and when HTTP/2 is not configured; controlled by RxDocumentClientImpl.useThinClient. - QueryPlan, metadata reads, and AllVersionsAndDeletes change feed continue to route through Compute Gateway (Gateway V1). - Probe failures are absorbed inside the orchestrator and the trigger is fire-and-forget on the GEM scheduler, so probe issues can never trip CosmosClient initialization or fail a topology refresh. - EndpointOrchestrator implements Closeable and is closed from GlobalEndpointManager.close() so no further probes are issued after client shutdown. THINCLIENT_ENABLED now defaults to true; opt out via COSMOS.THINCLIENT_ENABLED=false or COSMOS.THINCLIENT_PROBE_ENABLED=false. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Address PR #49437 review comments: probe single-flight, body-drain lifecycle, cancellable in-flight probes - EndpointOrchestrator: fold body-drain into probe Mono via flatMap+then(perProbeTimeout) so a slow/hanging response body cannot leak resources outside the cycle budget (Copilot #1, deep-review #3, jeet HIGH-2 minor). - EndpointOrchestrator: add single-flight CAS (cycleInProgress) plus monotonic cycle id; closed-check inside applyCycleResult drops late results so a post-close cycle cannot mutate health state (deep-review #1+#2, jeet HIGH-2). - EndpointOrchestrator: re-evaluate closed/feature-flag/endpoints at subscription time via Mono.defer so GEM.close() cancellation is honored before any HTTP I/O is issued. - GlobalEndpointManager: retain probe Disposable in AtomicReference; close() now disposes the in-flight probe subscription so probe work cannot outlive the GEM/CosmosClient (Copilot #2, deep-review #2). - CHANGELOG: moved entry to unreleased 4.82.0-beta.1, reworded to honestly describe optimistic startup, N=2 RED-to-fallback hysteresis, and Direct-mode/metadata exclusion (Copilot #3, deep-review #4). Tests: 45 unit tests pass (EndpointOrchestratorTests + ConfigsTests + ThinClientProbeWiringTests). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Address PR #49437 second-batch review feedback - LocationCache.getThinClientRegionalEndpoints now walks both read and write region endpoint maps so single-master write-region failures still flip the probe gate. - EndpointOrchestrator.forceUnhealthy(reason) provides a non-HTTP path to flip the gate; GlobalEndpointManager calls it when topology says thin-client is eligible but no regional endpoint resolves. - Symmetric hysteresis: new COSMOS.THINCLIENT_PROBE_RECOVERY_THRESHOLD (default 1) so operators can require N consecutive GREEN cycles before flipping back to proxy. - Extracted RxDocumentClientImpl.useThinClientStoreModel(...) body into package-private static shouldUseThinClientStoreModel for direct unit testability; added ThinClientRoutingGateTests covering 9 routing paths. - EndpointOrchestratorTests.stubResponse now returns Mono.empty() to avoid Unpooled.EMPTY_BUFFER refCnt underflow across multiple probe calls. - Removed unused locals; added recoveryThresholdRequiresMultipleGreenCycles, forceUnhealthy_flipsGateToRedWithoutRunningProbe, forceUnhealthy_onClosedOrchestrator_isNoOp tests. All 57 unit tests in the touched files pass. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Clarify CHANGELOG: probe recovery threshold is now configurable Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Address PR #49437 review (rounds 4-8): reactor-chain probe, FQN cleanup, fix gwV2Cto and ThinClient user-agent assertions - GlobalEndpointManager: convert thin-client probe trigger to a Mono<Void> chained into the topology-refresh reactor pipeline (replaces fire-and-forget subscribe). Removes thinClientProbeDisposable field and its close() handling since cancellation now propagates through the outer subscription. - EndpointProbeClient/EndpointProbeClientTests/ThinClientProbeWiringTests: replace inline FQNs with imports (java.io.Closeable, java.util.List, java.net.ConnectException, com.azure.cosmos.implementation.http.HttpHeaders). - ClientConfigDiagnosticsTest: compute gwV2Cto dynamically from Configs.isThinClientEnabled() so assertions remain valid after the default flip to true. - ConfigsTests: update default-threshold assertions from 2 to 1 to match DEFAULT_THINCLIENT_PROBE_FAILURE_THRESHOLD=1. - UserAgentContainerTest.UserAgentIntegration: expect '|F4' suffix because the ThinClient feature flag (1 << 2) is now included by default. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix CI test fallout from default ThinClient enablement UserAgentSuffixTest.validateUserAgentSuffix and CosmosDiagnosticsTest.generateHttp2OptedInUserAgentIfRequired: include UserAgentFeatureFlags.ThinClient in computed |F<hex> suffix when COSMOS.THINCLIENT_ENABLED is true (now default after Gateway V2 default enablement). Mirrors RxDocumentClientImpl.addUserAgentSuffix + UserAgentContainer.setFeatureEnabledFlagsAsSuffix behavior. SinglePartitionDocumentQueryTest.querySinglePartitionDocuments: spy on both gateway-proxy and thin-proxy and assert exactly one invocation. Previous code only spied on the proxy implied by useThinClient() config intent, which races with the probe-healthy gate -- routing AND's intent with isProxyProbeHealthy() so on first cycle the request may go through gateway even when thin-client is configured. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Disable thin-client probe in Http2PingKeepaliveTest The test installs an iptables DROP on thin-client port 10250 to verify that Http2PingHandler closes the broken connection after consecutive PING ACK timeouts and the recovery request uses a new connection on the same regional endpoint. After default Gateway V2 enablement, the connectivity probe also fires HTTP/2 POSTs to port 10250 on every account refresh. With iptables dropping that port, the probe trips proxyHealthy=false, useThinClient StoreModel() returns false, and the data plane request routes through Gateway V1 on port 443 -- which iptables is not dropping. Result: the PING handler never fires, the warm-up and recovery requests use the same gateway channel, and the assertion 'recovery channel must differ from initial' fails (both ended up as 77af2e47 on build 6419227). Set COSMOS.THINCLIENT_PROBE_ENABLED=false in beforeClass so the probe short-circuits to a no-op, EndpointProbeClient.proxyHealthy stays optimistically true, and the data plane request actually flows over port 10250 where the iptables DROP can take effect. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Disable thin-client probe by default for E2E tests in TestSuiteBase Build 6424287 surfaced two new failure patterns: 1. CosmosNotFoundTests.performBulkOnDeletedContainerWithGatewayV2 (45 failures) - asserts substatus 1003 from the thin-client routing path, but observed 0 because the data plane was routed to Gateway V1 instead of the proxy. 2. PerPartitionCircuitBreakerE2ETests.*Gateway (26 failures) - TestSuiteBase.assertThinClientEndpointUsed could not find any request whose endpoint contained ':10250/', i.e. nothing actually went to the thin-client proxy. Both patterns trace to the same source: the new connectivity probe is enabled by default, the proxy-side /connectivity-probe endpoint is not deployed in every CI test account yet, and the default failure threshold is 1. So after the first probe cycle the SDK marks the proxy unhealthy and routes data plane traffic to Gateway V1, which breaks tests that explicitly assert thin-client routing. Disable the probe by default in TestSuiteBase's static initializer (only when the property is not already set), so all E2E tests inherit deterministic, configuration-driven routing. Tests that exercise the probe itself (EndpointProbeClientTests, ThinClientProbeWiringTests) set the property explicitly in @BeforeMethod and are not affected. Also drop the now-redundant per-class override in Http2PingKeepaliveTest - the base class disables it, and the test's @afterclass clear would otherwise re-enable the probe for any subsequent E2E test sharing the JVM. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Revert "Disable thin-client probe by default for E2E tests in TestSuiteBase" This reverts commit ad9e3df. * Add per-class thinclient probe disable in CosmosNotFoundTests and PerPartitionCircuitBreakerE2ETests Companion to the prior revert. The revert undid the global TestSuiteBase probe disable (which masked production behaviour). This commit adds the necessary per-class disable to the two test classes whose assertions explicitly require thinclient routing: CosmosNotFoundTests (thinclient group) and PerPartitionCircuitBreakerE2ETests (fi-thinclient-multi-master group). Both clear the property in their @afterclass. Http2PingKeepaliveTest already has its own disable (restored by the revert). Production callers continue to get the connectivity probe ON by default with the production failure threshold. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Add unit tests for QueryPlan and stored-procedure thin-client routing Covers 5 new scenarios in ThinClientRoutingGateTests: - ExecuteStoredProcedure on a StoredProcedure resource routes to thin client - Non-execute StoredProcedure ops (Create) route to Gateway V1 - OperationType.QueryPlan routes to thin client - QueryPlan returns false when probe is unhealthy - ExecuteStoredProcedure returns false when probe is unhealthy Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Remove endpoint-probe content from QueryPlan PR branch Reverts the merge of jeet1995/thin-client-probe-flow that was brought into this PR earlier, while keeping the branch up-to-date with upstream/main. Net diff vs upstream/main is now only the QueryPlan and StoredProcedure Gateway V2 routing changes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Propagate hybrid query diagnostics Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Advertise CountIf query feature Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Address query plan review feedback Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Share thin client test property setup Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Clean up query range conversion state Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Thread DocumentCollection into readMany query-plan validation readManyByPartitionKeys validates any caller-supplied custom query by fetching a query plan and asserting it is single-partition and non-hybrid. Until now that validation called fetchQueryPlanForValidation with no DocumentCollection, so the plan request was pinned to Gateway V1 (useGatewayMode = (partitionKeyDefinition == null)). Thread the container's DocumentCollection from RxDocumentClientImpl.validateCustomQueryForReadManyByPartitionKeys -> DocumentQueryExecutionContextFactory.fetchQueryPlanForValidation so QueryPlanRetriever has the PartitionKeyDefinition needed to convert PartitionKeyInternal-formatted queryRanges from the thin client (Gateway V2) proxy into the EPK-hex Range<String> entries the query pipeline consumes. With this wiring, the validation query plan goes to the thin client when the client is configured for it, and remains on Gateway V1 otherwise. No behavior change on the non-thin-client path. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Cover ReadManyByPartitionKey QueryPlan bifurcation with unit + E2E tests - Add ReadManyByPartitionKeyQueryPlanRoutingTest unit tests that pin the useGatewayMode gate in QueryPlanRetriever: gateway mode when DocumentCollection is null and partitioned mode when a PartitionKeyDefinition is present. - Add three readManyByPartitionKeys E2E tests to ThinClientQueryE2ETest that exercise the validation QueryPlan path through Direct TCP (baseline) and Gateway V2 (thin client), covering no-custom-query, projection+filter, and parameterized variants. Each thin-client diagnostics page is asserted to use the :10250 endpoint. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Skip RNTBD QueryPlan frames when capturing V2 feed requests QueryPlan requests intentionally carry no RCS/CL headers (matches the V1 HTTP behavior). When the V2 thin-client routes the QueryPlan precursor through the same :10250 endpoint as the data query, the spy must skip the QueryPlan frame so the assertion checks the actual data-query frame. This mirrors the IS_QUERY_PLAN_REQUEST filter on the V1 path. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix thin-client QueryPlan error returning statusCode 0 instead of 400 The thin-client/Gateway-V2 proxy returns a non-2xx response with a raw, non-JSON, NUL-padded error body for invalid-syntax queries. In RxGatewayStoreModel.validateOrThrow, new CosmosError(body) attempted to parse that body as JSON and threw IllegalArgumentException, which escaped the method before the existing status-carrying throw could run. Upstream then wrapped it as statusCode 0. Wrap the CosmosError(body) construction in a narrow try/catch (IllegalArgumentException) and fall back to the non-parsing CosmosError(errorCode, message) constructor with a sanitized body. The existing throw now fires with the correct status (400) and the proxy error text is preserved. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Harden ThinClientQueryE2ETest assertions (F1-F5) F1: strict 400 + unconditional thin-client endpoint check on invalid query, locking the statusCode-0->400 fix. F2: ordered-vs-sorted-set document-ID comparison (ORDER BY sequence-compared, others set-compared). F3: ID-set equality + no-duplicate check across drained continuation pages. F4: numeric-tolerance (1e-6) comparison for scalar and GROUP BY aggregates to avoid float-formatting false mismatches. F5: validated vector/full-text/hybrid queries match Direct vs thin-client end-to-end through the proxy. Validated live via -Pthinclient: 84 tests, 0 failures. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Replace thin-client test matrix with reverse-engineered E2E test spec Rewrite THINCLIENT_TEST_MATRIX.md as a reviewable test-design specification reverse-engineered from the committed ThinClientQueryE2ETest code (84 tests). Documents the differential-testing oracle (Direct :443 baseline vs thin client :10250 SUT), the data-model fixture, every assertion contract (endpoint provenance, ordered-vs-unordered ID equality, scalar/GROUP BY tolerance), the full 84-test matrix, the F1-F5 hardened special cases (continuation draining, invalid-query 400, vector/FTS/hybrid ranking, readMany validation path), advertised-feature coverage, and known gaps (CountIf/DCount/MultipleOrderBy) for reviewer sign-off. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Apply Aditya Sarpotdar review feedback to thin-client E2E tests Executes actionable review feedback on ThinClientQueryE2ETest plus two harness fixes surfaced during live validation against thin-client-multi-region-ci. ThinClientQueryE2ETest: - Strict ordering: ORDER BY results validated with isStrictlyOrdered across the board (stricter-by-default, per reviewer guidance) instead of set parity. - Add testMultipleOrderBy() with a composite-index container to cover the MultipleOrderBy query feature. - Add testDCount() using the canonical Cosmos DCount idiom (COUNT over a DISTINCT VALUE subquery); SQL-standard COUNT(DISTINCT ...) is not valid Cosmos SQL grammar. - Reword multi-EPK-range comment/javadoc to reflect emulator/backend reality (multiple metadata ranges served by a single backend partition; SDK routing and query pipeline still exercised). Harness fixes: - TestSuiteBase: add wait-and-poll utility waitForCollectionToBeAvailableToRead (predicate on NotFound/substatus 1013) to deflake "Collection is not yet available for read" on freshly created containers; update call sites in OrderbyDocumentQueryTest, NonStreamingOrderByQueryVectorSearchTest, QueryValidationTests, ReadFeedCollectionsTest. - SinglePartitionDocumentQueryTest: make the processMessage Mockito assertion mode-aware (thin-client routes query plan + query => times(2)). Validated live against thin-client-multi-region-ci: 86/86 green. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Address review feedback: opt-in thin-client QueryPlan routing + kill-switch, stop advertising CountIf, remove PR-reference comments * Address review feedback: broaden gateway error fallback to ClassCastException, drop response ObjectNode mutation, fail-fast on empty queryRanges, align readMany routing test with opt-in * Make thin-client QueryPlan no-EPK-headers contract explicit in wrapInHttpRequest * Fix missing StandardCharsets import in RxGatewayStoreModelTest * Test: add QueryOracle-derived LIKE/scalar-expression thin-client parity tests; drain only first page in OperationPolicies readAllItems to avoid full-container enumeration timeout * Test: verify a cached proxy-generated query plan still executes correctly through the thin client * Test: cross-partition tests use a dedicated multi-physical-partition fixture and assert >1 partition key range contacted; enforce ordered full-row comparison when id is not projected; clearer naming * Add CHANGELOG entry for QueryPlan request routing to Gateway V2 * Call out Execute Stored Procedure support in CHANGELOG entry * Fix flaky region/timing-sensitive fault-injection E2E tests CustomerWorkflowPartitionLevelCircuitBreakerTest and CustomerWorkflowHighE2ETimeoutTest: bump the ThresholdBasedAvailabilityStrategy threshold from 100ms to 300ms so the cold primary-region connection can dispatch the request and record the injected fault before the availability-strategy hedge to a warm secondary region wins the race and cancels it (which previously yielded 0 fault hits). PerPartitionCircuitBreakerE2ETests: before asserting short-circuit routing, wait until the first preferred region is actually marked Unavailable, removing the race between the failure-counter estimate and the asynchronous routing exclusion. Scoped to single-faulty-region configs (expectedRegionCountWithFailures == 1) and run once so fault-in-all-regions configs (which can never short-circuit) do not hit the method timeout. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Remove test AAD auth toggle; move thin-client test matrix to docs/ and align with ThinClientQueryE2ETest * Add cross-partition aggregate / GROUP BY parity tests for thin-client query Adds 7 thin-client E2E tests covering cross-partition aggregate merge (COUNT/SUM/AVG/MIN/MAX) and GROUP BY, plus two helpers that assert Direct vs thin-client value parity and verify the query genuinely fanned out (distinctPartitionKeyRangesContacted > 1), exercising the cross-partition MERGE path the single-partition aggregate tests miss. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix thin-client MULTI_HASH prefix EPK over-span; add QueryPlan parity tests For a partial (prefix) hierarchical partition key, the thin-client / Gateway V2 path previously sent only StartEpkHash/EndEpkHash routing headers and no doc-level EPK filter, so the proxy resolved the request to the owning physical partition and returned every co-located document (e.g. 18 instead of 6). ThinClientStoreModel.wrapInHttpRequest now sets READ_FEED_KEY_TYPE= EffectivePartitionKeyRange, START_EPK and END_EPK on the request headers before RntbdRequest.from() so RntbdRequestHeaders.addStartAndEndKeys serializes the prefix EPK sub-range as the backend doc-level filter (hex string as UTF-8 bytes), mirroring the Direct/RNTBD FeedRangeEpkImpl path. StartEpkHash/EndEpkHash routing headers are retained. Adds ThinClientQueryE2ETest coverage: QueryPlan parity across sources and a hierarchical prefix half-open range test asserting thin-client matches Direct TCP. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Enforce QueryPlan thin-client routing invariant and harden container-readiness probe for thin-client-default-on - Remove the QueryPlan exemption in assertThinClientEndpointUsed: when ThinClient is opted in with HTTP/2 enabled, QueryPlan calls MUST also route to Gateway V2 (:10250), matching production default DEFAULT_THINCLIENT_QUERY_PLAN_ENABLED=true. - Restore public static visibility on ThinClientTestBase.assertThinClientEndpointUsed / assertGatewayEndpointUsed so cross-package static imports (CosmosMultiHashTest) compile. - Add OWNER_RESOURCE_NOT_EXISTS (subStatus 1003) to the retryable NOTFOUND allowlist in TestSuiteBase.isRetryableCollectionReadinessFailure. With thin-client default-on the beforeSuite readability probe (a QueryPlan) routes through GW V2; a freshly-created container not yet propagated to a secondary region returns 404/1003, which must be treated as a transient readiness failure instead of cascade-skipping the suite. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Add kill-switch flip validation for QueryPlan thin-client routing Validate the COSMOS.THINCLIENT_QUERY_PLAN_ENABLED sysprop/env kill-switch in both directions: - Unit: new case in ReadManyByPartitionKeyQueryPlanRoutingTest asserts QueryPlan is forced to Gateway V1 (useGatewayMode=true) when the flag is disabled, even with thin-client eligible. - E2E: make TestSuiteBase.assertThinClientEndpointUsed flag-aware so QueryPlan endpoint assertions flip based on the live-read flag while data operations continue to route to the thin-client endpoint. Both scenarios verified green (unit 3/3; E2E 112/112, oracle:true for flag ON and flag OFF). Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Restore ambient COSMOS.THINCLIENT_ENABLED in ThinClientTestBase teardown ThinClientTestBase teardown unconditionally cleared the COSMOS.THINCLIENT_ENABLED system property. Since Configs.isThinClientEnabled() reads the property lazily per client build, a ThinClientTestBase subclass running before a property-dependent class inherited from main (e.g. CosmosMultiHashTest, ThinClientQueryE2ETest) in the same JVM wiped the ambient -DCOSMOS.THINCLIENT_ENABLED=true flag supplied by the thin-client CI lane. Later clients then built with thin client disabled and routed to :443, breaking the :10250 thin-client endpoint assertions. Save the prior value on enable and restore it on teardown (clear only when it was originally absent), mirroring CosmosConsistencyOverrideValidationTest's restoreSystemProperty pattern. Helpers changed from static to instance so the saved value is per-instance, avoiding races across TestNG factory instances. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Remove redundant COSMOS.THINCLIENT_ENABLED mutation from ThinClientTestBase; revert workflow tests to main ThinClientTestBase no longer sets/clears COSMOS.THINCLIENT_ENABLED per class. The 'thinclient' Maven lane already sets THINCLIENT_ENABLED=true and HTTP2_ENABLED=true JVM-wide via failsafe systemPropertyVariables, and all thin classes run only in that lane (groups={thinclient}). The prior per-class clearProperty in @afterclass wiped the ambient -D for later byte-identical-to-main classes in the same JVM, mis-routing requests to :443. Relying on the lane flag (matching ThinClientQueryE2ETest) removes that footgun. Also revert CustomerWorkflowHighE2ETimeoutTest and CustomerWorkflowPartitionLevelCircuitBreakerTest to upstream/main (availability-strategy threshold stays 100ms), since this PR requires no changes there. Verified locally (AZURE_TEST_MODE=LIVE, -Pthinclient,query,consistency-overrides): 234 passed, 0 failures/skips; QueryPlan routes to :10250. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Abhijeet Mohanty <copilot@noreply.local>
Description
Adds a new readManyByPartitionKey API surface to the Java Cosmos SDK (sync + async) and wires it through the Spark connector to support PK-only reads (including partial HPK), with query-plan-based validation for custom queries.
Changes:
Added public readManyByPartitionKeys overloads in CosmosAsyncContainer / CosmosContainer and an internal AsyncDocumentClient + RxDocumentClientImpl implementation that groups PKs by physical partition and issues per-range queries.
Introduced ReadManyByPartitionKeyQueryHelper to compose PK filters into user-provided SQL and added a new config knob for per-partition batching.
Added Spark support (UDF + PK serialization/parsing helper + reader) and unit/integration tests for query composition and end-to-end behavior.
All SDK Contribution checklist:
General Guidelines and Best Practices
Testing Guidelines